package cn.edu.hit.kg;

import java.util.Arrays;
import java.util.Comparator;
import java.util.Objects;
import java.util.Random;
import java.util.function.BiPredicate;
import java.util.stream.Collectors;
import java.util.stream.IntStream;


public class SampleGenerator {
    private int featureCount, reasonCount;
    /**
     * 要求 各reason的概率和应当大于等于1 每个单独的概率应小于等于1  应当附加一个表示“未知原因”的原因
     */
    private double[] reasonProbabilities;
    private double[][] featureReasonProbabilities;

    public SampleGenerator(int featureCount, int reasonCount) {
        this.featureCount = featureCount;
        this.reasonCount = reasonCount;
        reasonProbabilities = new double[reasonCount];
        featureReasonProbabilities = new double[featureCount][reasonCount];
//        for (int i = 0; i < featureCount; i++) {
//            featureReasonProbabilities[i] = new double[reasonCount];
//        }
    }

    public SampleGenerator setRP(int reasonRank, double probability) {
        this.reasonProbabilities[reasonRank] = probability;
        return this;
    }

    public SampleGenerator setFP(int featureRank, int reasonRank, double probability) {
        this.featureReasonProbabilities[featureRank][reasonRank] = probability;
        return this;
    }

    public Integer[][] makeSamples(int sampleSize) {
        Integer[][] samples = new Integer[sampleSize][featureCount + reasonCount];
        for (int i = 0; i < samples.length; i++) {
            Arrays.fill(samples[i], 0);
        }
        int[] reasonPositions = new int[reasonCount];
        for (int i = 0; i < reasonCount; i++) {
//            int sumOfReason = Math.toIntExact(Math.round(reasonProbabilities[i] * sampleSize));
//            Row<Integer> column = new ColumnOfMartix<>(samples,featureCount+i);
//            randomSelect(column,sumOfReason,1,0);
            reasonPositions[i] = 1 << i;
        }

        for (int i = 0; i < featureCount; i++) {
//            BiPredicate<Integer,Integer> excludeCondition = (index,val)-> false;
            int[] distribute = makeDistribution(sampleSize, i);
            int[] sampleDistribute = new int[sampleSize];
            Row<Integer> sampleDistributeRow = new Row<Integer>() {
                @Override
                public int size() {
                    return sampleDistribute.length;
                }

                @Override
                public Integer get(int index) {
                    return sampleDistribute[index];
                }

                @Override
                public void set(int index, Integer value) {
                    sampleDistribute[index] = value;
                }
            };
            for (int k = 1; k < distribute.length; k++) {
                Row<Integer> row = new ColumnWithFilter<>(sampleDistributeRow, (index, val) -> val == 0);
                randomSelect(row, distribute[k], k, 0);
            }

            for (int m = 0; m < sampleDistribute.length; m++) {
                if (sampleDistribute[m] == 0)
                    continue;
                for (int n = 0; n < reasonPositions.length; n++) {
                    if ((sampleDistribute[m] & reasonPositions[n]) > 0) {
                        samples[m][i] = 1;
                        samples[m][featureCount + n] = 1;
                    }
                }
            }

            //先计算reason之间 公用的feature样本数
//            for (int j = 0; j < reasonCount; j++) {
//                int sum = Math.toIntExact(Math.round(featureReasonProbabilities[i][j] *reasonProbabilities[j]* sampleSize));
//                int finalJ = j;
//                int selectedSum = new ColumnWithFilter<>(new ColumnOfMartix<>(samples,i),((index, val) -> samples[index][featureCount+ finalJ]==1 && val==1 )).size();
//                System.out.printf("(%d,%d)---->(%d,%d)\n",i,j,sum,selectedSum);
//                BiPredicate<Integer, Integer> filterBiPredicate = (index, val) -> samples[index][featureCount + finalJ] == 1 && val == 0;
//                Row<Integer> columnForReason = new ColumnWithFilter<>(new ColumnOfMartix<>(samples,i), filterBiPredicate/*.and(excludeCondition.negate())*/);
//                randomSelect(columnForReason,sum-selectedSum,1,0);
////                excludeCondition = excludeCondition.or(filterBiPredicate);
//            }
        }
        return samples;
    }

    /**
     * 生成一个指定的分布
     *
     * @param sampleSize 样本总数
     * @param splitSizeP 存在splitSizeP.length个属性，每个样本可以包含一个或多个属性，数组中的每个元素代表该属性占全部样本的比例
     * @return 长度为 1 << splitSizeP.length 的int数组,是每个样本携带属性数最少的一个分布。数组的第i个元素存储该类样本的数量，对于 i&(1<<x) 成立即代表含有第x个属性
     */
    public static int[] makeDistribution(int sampleSize, double[] splitSizeP) {
        int[] distribution = new int[(1 << splitSizeP.length)];
        Requirement[] reasonRequires = new Requirement[splitSizeP.length];
        int sumOfRequires = 0;
        for (int i = 0; i < splitSizeP.length; i++) {
            int reasonRequireSize = Math.toIntExact(Math.round(sampleSize * splitSizeP[i]));
            sumOfRequires += reasonRequireSize;
            reasonRequires[i] = new Requirement(i, reasonRequireSize);
        }
        //计算得到需要生成的重叠数
        int rounds = (int) Math.ceil(sumOfRequires * 1.0 / sampleSize);
//        System.out.println("Rounds:"+rounds);
        Arrays.sort(reasonRequires, Comparator.comparingInt(x -> x.requireSize));
//        System.out.println(Arrays.toString(reasonRequires));

        for (int i = 0; ; i++) {
            int leftAmount = Math.min(sampleSize, Arrays.stream(reasonRequires).mapToInt(Requirement::getRequireSize).sum());
            if (leftAmount == 0 || i > splitSizeP.length) {
                break;
            }
            for (int i1 = 0; i1 < reasonRequires.length; i1++) {
                Requirement require = reasonRequires[i1];
                if (leftAmount == 0 || require.requireSize == 0) {
                    continue;
                }
//                System.out.println("当前需求："+require);
//                System.out.println("当前剩余："+leftAmount);
                int pieceOfDivide = leftAmount / (splitSizeP.length - i1);
//                System.out.println(pieceOfDivide);
                int afford = Math.min(require.requireSize, pieceOfDivide);

//                System.out.println(Arrays.toString(distribution));
                int realAfford = distributeAffords(distribution, splitSizeP.length, require.rank, afford, i);//将afford分发到各个可重叠的部分
//                System.out.println(Arrays.toString(distribution));
//                System.out.println(afford);
//                System.out.println(realAfford);
                require.requireSize -= realAfford;
                leftAmount -= realAfford;
            }
        }
        return distribution;
    }

    public int[] makeDistribution(int sampleSize, int featureIndex) {
        double[] splitSizeP = IntStream.range(0, reasonCount)
                .mapToDouble(i -> reasonProbabilities[i] * featureReasonProbabilities[featureIndex][i])
                .toArray();
        return makeDistribution(sampleSize, splitSizeP);
    }

    /**
     * 生成一个指定的二维分布
     *
     * @param sampleSize     样本总数
     * @param rowRequires    第一维的属性，每个样本可以包含一个或多个第一维的属性，数组中的每个元素代表该属性占全部样本的比例
     * @param columnRequires 第二维的属性，每个样本可以包含一个或多个第二维的属性，数组中的每个元素代表该属性占全部样本的比例
     * @param matrix         组合matrix[x][y] 即同时含有x,y两个属性的样本占全部样本的比例
     * @return 返回一个1 << (rowRequires.length + columnRequires.length)的int数组,是每个样本携带属性数最少的一个分布。
     */
    public static int[] makeDistributions(int sampleSize, double[] rowRequires, double[] columnRequires, double[][] matrix) {
        //先对column做拆分
        int rv2 = 1 << columnRequires.length;
        // 每row做一下拆分，最后一行是合计
        int[][] d_row = new int[rowRequires.length + 1][rv2];
        for (int i = 0; i < rowRequires.length; i++) {
            int[] distribution = makeDistribution(Math.toIntExact(Math.round(sampleSize * rowRequires[i])), matrix[i]);
            System.arraycopy(distribution, 0, d_row[i], 0, distribution.length);
            for (int j = 1; j < rv2; j++) {
                d_row[rowRequires.length][j] += distribution[j];
            }
        }

        //再对row做拆分 TODO 此处使用 columnRequires 计算得到的与上面的不一致，应改为由上面结果得到
        int rd2 = 1 << rowRequires.length;
        int[][] d_column = new int[rv2 + 1][rd2];
        for (int i = 1; i < rv2; i++) {
            int finalI = i;
            double[] splitSizeP = IntStream.range(0, rowRequires.length).mapToDouble(row -> d_row[row][finalI] * 1.0 / sampleSize).toArray();
            System.out.println(Arrays.toString(splitSizeP));
//            int[] distribution = makeDistribution(d_row[rowRequires.length][i], splitSizeP);
            int[] distribution = makeDistribution(sampleSize, splitSizeP);
            System.arraycopy(distribution, 0, d_column[i - 1], 0, distribution.length);
            for (int j = 1; j < rd2; j++) {
                d_column[rv2][j] += distribution[j];
            }
        }

//        TODO 综合上面两个方向的拆分结果，生成最终的方式：
        System.out.println(Arrays.deepToString(d_row));
        System.out.println(Arrays.deepToString(d_column));

        return null;
    }

    private static int distributeAffords(int[] distribution, int reasonCount, int reasonRank, int afford, int round) {
        if (round == 0) {
            distribution[1 << reasonRank] = afford;
            return afford;
        } else {
            int combinations = combination(reasonCount - round - 1, reasonCount - 1);
            Requirement[] requirements = new Requirement[combinations];
            int reasonPosit = 1 << reasonRank;

            for (int i = 1, index = 0; i <= (1 << reasonCount) && index < combinations; i++) {
                if ((i & reasonPosit) > 0 || countBits(i) != round)
                    continue;
                requirements[index++] = new Requirement(i, distribution[i]);
            }

//            System.out.println(Arrays.toString(requirements));

            int arranged = 0;//本次分配的总数
            Arrays.sort(requirements, Comparator.comparingInt(x -> x.requireSize));
//            System.out.println("具体分配中的需求："+ Arrays.toString(requirements));
            for (int i = 0; i < combinations; i++) {
                Requirement requirement = requirements[i];
                if (requirement.requireSize == 0)
                    continue;
                int divide = Math.min(requirement.requireSize, afford / (combinations - i));
                requirement.requireSize -= divide;
                distribution[requirement.rank] -= divide;
                distribution[requirement.rank | reasonPosit] += divide;
                afford -= divide;
                arranged += divide;
            }
            return arranged;
        }
    }

    /**
     * 计算二进制数中包含几个1
     *
     * @param n
     * @return
     */
    static int countBits(int n) {
        int count = 0;
        while (n != 0) {
            n = n & (n - 1);
            count++;
        }
        return count;
    }

    /**
     * 排列组合问题求组合数
     *
     * @param groupSize 每个分组的元素个数
     * @param totalSize 所有元素的总个数
     * @return 分组总数
     */
    static int combination(int groupSize, int totalSize) {
        if (groupSize < totalSize - groupSize) {
            return combination(totalSize - groupSize, totalSize);
        }
        long amount = 1;
        for (int i = totalSize; i > groupSize; i--) {
            amount *= i;
        }
        for (int i = 1; i <= totalSize - groupSize; i++) {
            amount /= i;
        }
        return Math.toIntExact(amount);
    }

    static class Requirement {
        public final int rank;
        public int requireSize;

        public Requirement(int rank, int requireSize) {
            this.rank = rank;
            this.requireSize = requireSize;
        }

        @Override
        public String toString() {
            return "(" + rank +
                    ", " + requireSize +
                    ')';
        }

        public int getRequireSize() {
            return requireSize;
        }
    }

    public static void main(String[] args) {
//        try (NeoDriver driver = new NeoDriver()){
//            driver.executeVoid(session -> {
//                session.run("match (n:Fault) return n").list(x->x.get("n")).forEach(f->{
//                    System.out.println(f.get("name").asString());
//                    System.out.println(f.asMap());
//                    List<Double> matrix = f.get("matrix").asList(x->x.asDouble());
//                    int fn = matrix.get(0).intValue();
//                    int rn = matrix.get(1).intValue();
//                    SampleGenerator generator = new SampleGenerator(fn,rn);
//                    for (int i = 0; i < rn; i++) {
//                        System.out.println("r"+i+matrix.get((rn+1)*fn+i+2));
//                        generator.setRP(i,matrix.get((rn+1)*fn+i+2));
//                    }
//
//                    for (int i = 0; i < fn; i++) {
//                        for (int j = 0; j < rn; j++) {
//                            System.out.println("fp:"+matrix.get((rn+1)*i+j+2));
//                            generator.setFP(i,j,matrix.get((rn+1)*i+j+2));
//                        }
//                    }
//
//                    Integer[][] samples = generator.makeSamples(10000);
//                    System.out.println("f1\tf2\tr1\tr2\tr3");
//                    int[] sums = new int[5];
//                    for (Integer[] sample : samples) {
//                        System.out.printf("%d\t%d\t%d\t%d\t%d\n",sample[0],sample[1],sample[2],sample[3],sample[4]);
//                        for (int i = 0; i < (5); i++) {
//                            sums[i] += sample[i];
//                        }
//                    }
//                    System.out.println(Arrays.toString(sums));
//                });
//            });
//        }
//
//
//        System.exit(0);

//        int[] d0 = {0, 50, 400, 350, 25, 0, 175, 0};
//        System.out.println(distributeAffords(d0, 3, 0, 450, 1));
//        System.out.println(Arrays.toString(d0));


        double[] rowRequire = new double[]{0.85, 0.75, 0.2};
        double[] columnRequire = new double[]{0.65, 0.75};
        double[][] matrix = new double[][]{
                {0.8, 0.9},
                {0.75, 0.65},
                {0.45, 0.95}
        };
        makeDistributions(1000, rowRequire, columnRequire, matrix);

        System.exit(0);

        int[] d1 = makeDistribution(1000, new double[]{0.85, 0.75, 0.2});
        System.out.println(IntStream.range(0, d1.length).mapToObj(Integer::toBinaryString).mapToInt(Integer::parseInt)
                .mapToObj(i -> String.format("%03d", i)).collect(Collectors.joining("\t")));
        System.out.println(Arrays.stream(d1).mapToObj(Objects::toString).collect(Collectors.joining("\t")));
        System.exit(0);

        SampleGenerator generator = new SampleGenerator(3, 2);
        generator.setRP(0, 0.85)
                .setRP(1, 0.75)
                .setFP(0, 0, 0.70)
                .setFP(0, 1, 0.10)
                .setFP(1, 0, 0.15)
                .setFP(1, 1, 0.80)
                .setFP(2, 0, 0.75)
                .setFP(2, 1, 0.60);

        System.out.println(Arrays.toString(generator.makeDistribution(100, 0)));
        System.out.println(Arrays.toString(generator.makeDistribution(100, 1)));
        System.out.println(Arrays.toString(generator.makeDistribution(100, 2)));
        Integer[][] samples = generator.makeSamples(10000);
        System.out.println("f1\tf2\tf3\tr1\tr2");
        int[] sums = new int[5];
        for (Integer[] sample : samples) {
            System.out.printf("%d\t%d\t%d\t%d\t%d\n", sample[0], sample[1], sample[2], sample[3], sample[4]);
            for (int i = 0; i < (5); i++) {
                sums[i] += sample[i];
            }
        }
        System.out.println(Arrays.toString(sums));
    }

    interface Row<T> {
        int size();

        T get(int index);

        void set(int index, T value);

        default void fill(T value) {
            for (int i = 0; i < size(); i++) {
                set(i, value);
            }
        }
    }

    static class ColumnOfMartix<T> implements Row<T> {
        protected T[][] data;
        protected int column;

        public ColumnOfMartix(T[][] data, int column) {
            this.data = data;
            this.column = column;
        }

        @Override
        public int size() {
            return data.length;
        }

        @Override
        public T get(int index) {
            return data[index][column];
        }

        @Override
        public void set(int index, T value) {
            data[index][column] = value;
        }
    }

    static class ColumnWithFilter<T> implements Row<T> {
        private Row<T> superRow;
        private BiPredicate<Integer, T> rowCondition;
        private int initialSize = 0;
        private int[] indexMap;

        public ColumnWithFilter(Row<T> superRow, BiPredicate<Integer, T> rowCondition) {
            this.superRow = superRow;
            this.rowCondition = rowCondition;
            int[] tempIndexes = new int[superRow.size()];
            for (int i = 0; i < superRow.size(); i++) {
                if (rowCondition.test(i, superRow.get(i))) {
                    tempIndexes[initialSize] = i;
                    initialSize++;
                }
            }
            indexMap = new int[initialSize];
            System.arraycopy(tempIndexes, 0, indexMap, 0, initialSize);
        }

        @Override
        public int size() {
            return initialSize;
        }

        @Override
        public T get(int index) {
            return superRow.get(indexMap[index]);
        }

        @Override
        public void set(int index, T value) {
            superRow.set(indexMap[index], value);
        }
    }

    public static <T> void randomSelect(Row<T> row, int randomSize, T valueToSet, T initialValue) {
        int size = row.size();
        if (randomSize >= size) {
            row.fill(valueToSet);
            return;
        }
        if (randomSize > size / 2) {
            row.fill(valueToSet);
            randomSelect(row, size - randomSize, initialValue, valueToSet);
            return;
        }
        Random random = new Random();
        for (int i = 0; i < randomSize; i++) {
            int selectedIndex = random.nextInt(size - i);
            while (valueToSet.equals(row.get(selectedIndex)) && selectedIndex < size * 2) {
                selectedIndex++;
            }
            row.set(selectedIndex % row.size(), valueToSet);
        }
    }
}
